# !/usr/bin/python
# -*- coding: utf-8 -*-
import numpy as np
import pandas as pd
import sklearn
from resnet import Classifier_RESNET
import mlflow


def creat_train_test(dataset_path, fault_type):
    data = pd.read_csv(dataset_path)
    data = data.sample(frac=1).reset_index(drop=True)
    train_data = []
    test_data = []
    for i in range(len(data)):
        temp_data = list(data.iloc[i, 1:4])
        temp_data.append(data[fault_type][i])
        if i % 2 != 0:
            train_data.append(temp_data)
        else:
            test_data.append(temp_data)

    train_data = np.array(train_data)
    test_data = np.array(test_data)
    return train_data, test_data


def train_test_set(datasets_path, fault_type):
    train_data, test_data = creat_train_test(datasets_path, fault_type)

    train_data = np.array(train_data)
    x_train = np.array(train_data[:, :-1])
    y_train = train_data[:, -1]

    test_data = np.array(test_data)
    x_test = np.array(test_data[:, :-1])
    y_test = test_data[:, -1]

    return x_train, y_train, x_test, y_test


def fit_classifier(datasets_path, fault_type, Conv_num, batch_size, nb_epochs, output_directory):
    # 模型训练
    x_train, y_train, x_test, y_test = train_test_set(datasets_path, fault_type)
    nb_classes = len(np.unique(np.concatenate((y_train, y_test), axis=0)))

    # transform the labels from integers to one hot vectors
    enc = sklearn.preprocessing.OneHotEncoder(categories='auto')
    enc.fit(np.concatenate((y_train, y_test), axis=0).reshape(-1, 1))
    y_train = enc.transform(y_train.reshape(-1, 1)).toarray()
    y_test = enc.transform(y_test.reshape(-1, 1)).toarray()
    y_true = np.argmax(y_test, axis=1)

    if len(x_train.shape) == 2:
        # add a dimension to make it multivariate with one dimension
        x_train = x_train.reshape((x_train.shape[0], x_train.shape[1], 1))
        x_test = x_test.reshape((x_test.shape[0], x_test.shape[1], 1))

    input_shape = x_train.shape[1:]
    print("创建模型")
    classifier = create_classifier('resnet', input_shape, nb_classes, Conv_num, output_directory, verbose=True)
    print("模型训练")
    classifier = classifier.fit(x_train, y_train, x_test, y_test, y_true, batch_size, nb_epochs)
    mlflow.keras.log_model(classifier.model_best, "model")


def create_classifier(classifier_name, input_shape, output_shape, Conv_num, output_directory, verbose=False):
    if classifier_name == 'resnet':
        return Classifier_RESNET(output_directory, input_shape, output_shape, Conv_num, verbose)

